package geppetto.cat.alignments.phrasal;

import java.io.BufferedReader;
import java.io.DataOutputStream;
import java.io.File;
import java.io.FileOutputStream;
import java.io.FileReader;
import java.io.IOException;
import java.io.PrintStream;
import java.util.HashMap;
import java.util.zip.GZIPOutputStream;

import geppetto.phraseHMM.DoubleDistortionTable;
import geppetto.phraseHMM.phraseExtraction.extractedphrase.ExtractedPhrasePair;
import geppetto.reordering.graph.ReorderingGraph;
import geppetto.reordering.graph.ReorderingGraphEdge;

public class PhraseDistortionTable {
	public DoubleDistortionTable distortionTable;
	public DoubleDistortionTable distortionCount;
	double phrasePenalty = 0.9;
	double segmentationPenalty = 1; // no penalty needed because of agreement

	public PhraseDistortionTable() {
		super();
	}

	public double getTransition(ReorderingGraphEdge edge){
		return getTransition(edge.getToNode(), edge.getFromNode());
	}

	public double getTransition(ExtractedPhrasePair toNode, ExtractedPhrasePair fromNode){
		if(ReorderingGraph.isStartNode(fromNode)){		
			return distortionTable.getFirstProb(toNode.getSourceStartPos(), toNode.getSourceSentence().length) * Math.pow(phrasePenalty, toNode.getSourcePhrase().length) * Math.pow(segmentationPenalty, toNode.getTargetPhrase().length -1);
		}
		if(ReorderingGraph.isEndNode(toNode)){
			int sentenceLength = toNode.getSourceSentence().length;
			//System.out.println(toNode.getSourceStartPos() + " " + fromNode.getSourceEndPos() % sentenceLength + " " + sentenceLength);
			return distortionTable.getDistProb(sentenceLength - 1, fromNode.getSourceEndPos() % sentenceLength, sentenceLength);
		}
		int sentenceLength = fromNode.getSourceSentence().length;
		if(toNode.isNullPhrase()){
			return distortionTable.getNullDistProb(toNode.getSourceStartPos() % sentenceLength, fromNode.getSourceStartPos() % sentenceLength, sentenceLength);
		}
		return distortionTable.getDistProb(toNode.getSourceStartPos(), fromNode.getSourceEndPos() % sentenceLength, sentenceLength) *
		Math.pow(phrasePenalty, toNode.getSourcePhrase().length) *
		Math.pow(segmentationPenalty, toNode.getTargetPhrase().length - fromNode.getTargetPhrase().length - 1);
	}

	public void addTransitionExpectedCount(ReorderingGraphEdge edge, double posterior){
		addTransitionExpectedCount(edge.getToNode(), edge.getFromNode(), posterior);
	}
	
	public void addTransitionExpectedCount(DoubleDistortionTable table, ReorderingGraphEdge edge, double posterior){
		addTransitionExpectedCount(table, edge.getToNode(), edge.getFromNode(), posterior);
	}

	public void addTransitionExpectedCount(ExtractedPhrasePair toNode, ExtractedPhrasePair fromNode, double posterior){
		addTransitionExpectedCount(distortionCount, toNode, fromNode, posterior);
	}

	public void addTransitionExpectedCount( DoubleDistortionTable table, ExtractedPhrasePair toNode, ExtractedPhrasePair fromNode, double posterior){
		if(ReorderingGraph.isStartNode(fromNode)){		
			table.addToFirst(posterior, toNode.getSourceStartPos());
		}
		else 
			if(ReorderingGraph.isEndNode(toNode)){
				int sentenceLength = toNode.getSourceSentence().length;
				table.addToCount(posterior, fromNode.getSourceEndPos() % sentenceLength, sentenceLength - 1);
			}
			else{
				int sentenceLength = fromNode.getSourceSentence().length;
				if(toNode.isNullPhrase()){
					table.addToNullCount(posterior, fromNode.getSourceStartPos(), toNode.getSourceStartPos() % sentenceLength);
				}
				else{
					table.addToCount(posterior, fromNode.getSourceEndPos() % sentenceLength, toNode.getSourceStartPos());
				}
			}
	}

	public void loadDistortionTableFromDir(String directory){
		distortionTable = new DoubleDistortionTable(directory);
		distortionCount = new DoubleDistortionTable(directory);
	}

	public void printExtractionTable(PrintStream out, ReorderingGraphEdge[] edges){
		DoubleDistortionTable emptyDistortionTable = new DoubleDistortionTable(distortionTable.distProbs.length,0);
		for(ReorderingGraphEdge edge : edges){
			addTransitionExpectedCount(emptyDistortionTable, edge, edge.getPotential());
		}
		for(int i=0; i< emptyDistortionTable.distProbs.length; i++){
			out.println(i + " ||| " + emptyDistortionTable.distProbs[i] + " " + emptyDistortionTable.firstProbs[i] + " " + emptyDistortionTable.lastProbs[i]);
		}
		out.println("+++ ||| " + emptyDistortionTable.positiveOverflow);
		out.println("--- ||| " + emptyDistortionTable.negativeOverflow);
	}
	
	public void printExtractionTable(PrintStream out){
		for(int i=0; i< distortionCount.distProbs.length; i++){
			out.println(i + " ||| " + distortionCount.distProbs[i] + " " + distortionCount.firstProbs[i] + " " + distortionCount.lastProbs[i]);
		}
		out.println("+++ ||| " + distortionCount.positiveOverflow);
		out.println("--- ||| " + distortionCount.negativeOverflow);
	}
	
	public void printReorderingTable(PrintStream out){
		for(int i=0; i< distortionTable.distProbs.length; i++){
			out.println(i + " ||| " + distortionTable.distProbs[i] + " " + distortionTable.firstProbs[i] + " " + distortionTable.lastProbs[i]);
		}
		out.println("+++ ||| " + distortionTable.positiveOverflow);
		out.println("--- ||| " + distortionTable.negativeOverflow);
	}

	public void clearCountTables(){
		distortionCount.clear();
	}

	public void updateDistortionTableScores() {
		((DoubleDistortionTable)distortionTable).copy((DoubleDistortionTable)distortionCount);
	}

	public void loadCountsFromExtraction(String reorderingExtractFile) throws NumberFormatException, IOException {
		BufferedReader reorderingExtractReader = new BufferedReader(new FileReader(new File(reorderingExtractFile)));
		HashMap<Integer, Double> firstProbs = new HashMap<Integer, Double>();
		HashMap<Integer, Double> lastProbs = new HashMap<Integer, Double>();
		HashMap<Integer, Double> probs = new HashMap<Integer, Double>();
		double positiveOverflow = 0;
		double negativeOverflow = 0;
		int maxLen = 0;
		while(reorderingExtractReader.ready()){
			String line = reorderingExtractReader.readLine();
			String[] lineArray = line.split("\\s+\\|\\|\\|\\s+");
			String distortion = lineArray[0];
			if(distortion.equals("+++")){
				double prob = Double.parseDouble(lineArray[1]);
				positiveOverflow = prob;
			}
			else if(distortion.equals("---")){
				double prob = Double.parseDouble(lineArray[1]);
				negativeOverflow = prob;
			}
			else{
				int distance = Integer.parseInt(distortion);
				String[] features = lineArray[1].split("\\s+");
				Double probsDistance = Double.parseDouble(features[0]);
				Double probsFirst = Double.parseDouble(features[1]);
				Double probsLast = Double.parseDouble(features[2]);
				probs.put(distance, probsDistance);
				firstProbs.put(distance, probsFirst);
				lastProbs.put(distance, probsLast);
				maxLen=Math.max(maxLen, distance);
			}
		}
		double[] distProbsArray = new double[maxLen+1];
		double[] firstProbsArray = new double[maxLen+1];
		double[] lastProbsArray = new double[maxLen+1];		
		for(int key : firstProbs.keySet()){
			distProbsArray[key] = probs.get(key);
			lastProbsArray[key] = lastProbs.get(key);
			firstProbsArray[key] = firstProbs.get(key);
		}
		distortionCount = new DoubleDistortionTable(firstProbsArray, lastProbsArray, distProbsArray, positiveOverflow, negativeOverflow);
		distortionTable = new DoubleDistortionTable(maxLen+1,0);
	}

	public void loadDistortionTableFromFile(
			String initialDistortionTableDirectory) throws NumberFormatException, IOException {
		BufferedReader reorderingExtractReader = new BufferedReader(new FileReader(new File(initialDistortionTableDirectory)));
		HashMap<Integer, Double> firstProbs = new HashMap<Integer, Double>();
		HashMap<Integer, Double> lastProbs = new HashMap<Integer, Double>();
		HashMap<Integer, Double> probs = new HashMap<Integer, Double>();
		double positiveOverflow = 0;
		double negativeOverflow = 0;
		int maxLen = 0;
		while(reorderingExtractReader.ready()){
			String line = reorderingExtractReader.readLine();
			String[] lineArray = line.split("\\s+\\|\\|\\|\\s+");
			String distortion = lineArray[0];
			if(distortion.equals("+++")){
				double prob = Double.parseDouble(lineArray[1]);
				positiveOverflow = prob;
			}
			else if(distortion.equals("---")){
				double prob = Double.parseDouble(lineArray[1]);
				negativeOverflow = prob;
			}
			else{
				int distance = Integer.parseInt(distortion);
				String[] features = lineArray[1].split("\\s+");
				Double probsDistance = Double.parseDouble(features[0]);
				Double probsFirst = Double.parseDouble(features[1]);
				Double probsLast = Double.parseDouble(features[2]);
				probs.put(distance, probsDistance);
				firstProbs.put(distance, probsFirst);
				lastProbs.put(distance, probsLast);
				maxLen=Math.max(maxLen, distance);
			}
		}
		double[] distProbsArray = new double[maxLen+1];
		double[] firstProbsArray = new double[maxLen+1];
		double[] lastProbsArray = new double[maxLen+1];		
		for(int key : firstProbs.keySet()){
			distProbsArray[key] = probs.get(key);
			lastProbsArray[key] = lastProbs.get(key);
			firstProbsArray[key] = firstProbs.get(key);
		}
		distortionTable = new DoubleDistortionTable(firstProbsArray, lastProbsArray, distProbsArray, positiveOverflow, negativeOverflow);
		distortionCount = new DoubleDistortionTable(maxLen+1,0);
	}

}
